'''
Test utilities
'''
from contextlib import contextmanager
import datetime
from importlib import import_module

import pymysql


class EndAlert(Exception):
    '''End of production'''


class ProducerCounter():
    '''Producer only counting'''

    def __init__(self):
        self.index = 0

    def send(self, _topic, _data):
        '''Mock of send method, counting message'''
        self.index += 1


def get_now():
    '''
    Get current datetime.

    Time zone shall be consistent with database
    '''
    return datetime.datetime.utcnow()


@contextmanager
def _auto_commit_db():
    '''Prepare database for insert'''
    conn = pymysql.connect(**import_module('config_test').DB_CONFIG)
    yield conn
    conn.commit()
    conn.close()


SQL_DATA = 'INSERT INTO test (name, timestamp) VALUES (%s, %s)'


def insert_test_data(name, timestamp):
    '''Create new test data'''
    with _auto_commit_db() as conn:
        with conn.cursor() as cursor:
            cursor.execute(SQL_DATA, (name, timestamp))


def insert_time_drift(drift, applied_at):
    '''Create new time_drift'''
    with _auto_commit_db() as conn:
        with conn.cursor() as cursor:
            cursor.execute('INSERT INTO time_drift (duration_ms, applied_at)'
                           ' VALUES (%s, %s)', (drift, applied_at))
